Hi @will-cromar , could you please review? Thanks!

Would it be possible for you to move the plugin to your own repository? The packages in this repository are the ones that we (ie the PyTorch/XLA team) test, release, and maintain ourselves. Additional device support should be maintained out of tree.

You should be able to build the actual plugin directly from OpenXLA without having to create your own bazel workspace (something like bazel build --config=rocm //xla/pjrt/c:pjrt_c_api_gpu_plugin.so). Most of the mess of bazel build config in this repository supports the main torch_xla package build and won't be relevant for the plugin.

The Python parts of the package all look right to me.

We can accept the change in __init__.py since ROCm is a special case. We can also include a link to your plugin in our README.

Thanks for the feedback! I will have to further discuss with my team about setting up our own repository, and will reach back once it is up. :)